package cn.saltybe.analysis.crawler;

import java.io.*;
import java.net.URL;
import java.util.ArrayList;
import java.util.List;
import java.util.Map;

import com.geccocrawler.gecco.downloader.AbstractDownloader;
import com.geccocrawler.gecco.downloader.UserAgent;
import org.apache.http.HttpHost;

import com.gargoylesoftware.htmlunit.BrowserVersion;
import com.gargoylesoftware.htmlunit.HttpMethod;
import com.gargoylesoftware.htmlunit.WebClient;
import com.gargoylesoftware.htmlunit.WebRequest;
import com.gargoylesoftware.htmlunit.WebResponse;
import com.gargoylesoftware.htmlunit.html.HtmlPage;
import com.gargoylesoftware.htmlunit.util.NameValuePair;
import com.geccocrawler.gecco.downloader.DownloadException;
import com.geccocrawler.gecco.request.HttpPostRequest;
import com.geccocrawler.gecco.request.HttpRequest;
import com.geccocrawler.gecco.response.HttpResponse;
import com.geccocrawler.gecco.spider.SpiderThreadLocal;
import com.geccocrawler.gecco.utils.UrlUtils;

/**
 * 利用htmlunit实现js的dom操作和ajax
 *
 * @author huchengyi
 */
@com.geccocrawler.gecco.annotation.Downloader("htmlUnitDownloader")
public class HtmlUnitDownloader extends AbstractDownloader {

    private WebClient webClient;

    public HtmlUnitDownloader() {
    }

    public void start() {
        this.webClient = new WebClient(BrowserVersion.CHROME);
        this.webClient.getOptions().setPrintContentOnFailingStatusCode(false);
        this.webClient.getOptions().setThrowExceptionOnFailingStatusCode(false);
        this.webClient.getOptions().setThrowExceptionOnScriptError(false);
        this.webClient.getOptions().setRedirectEnabled(true);
        this.webClient.getOptions().setCssEnabled(false);
        this.webClient.setJavaScriptTimeout(3000);
        this.webClient.setJavaScriptErrorListener(new GeccoJavaScriptErrorListener());
    }

    public HttpResponse download(HttpRequest request, int timeout) throws DownloadException {
        System.out.println("Init Web Client...");
        start();
        System.out.println("Start download page: " + request.getUrl());
        try {
            URL url = new URL(request.getUrl());
            WebRequest webRequest = new WebRequest(url);
            webRequest.setHttpMethod(HttpMethod.GET);
            if (request instanceof HttpPostRequest) {//post
                HttpPostRequest post = (HttpPostRequest) request;
                webRequest.setHttpMethod(HttpMethod.POST);
                List<NameValuePair> requestParameters = new ArrayList<>();
                for (Map.Entry<String, String> entry : post.getFields().entrySet()) {
                    NameValuePair nvp = new NameValuePair(entry.getKey(), entry.getValue());
                    requestParameters.add(nvp);
                }
                webRequest.setRequestParameters(requestParameters);
            }
            //header
            boolean isMobile = SpiderThreadLocal.get().getEngine().isMobile();
            webRequest.setAdditionalHeader("User-Agent", UserAgent.getUserAgent(isMobile));
            webRequest.setAdditionalHeaders(request.getHeaders());
            //proxy
            HttpHost proxy = SpiderThreadLocal.get().getEngine().getProxysLoader().getProxy();
            if (proxy != null) {
                webRequest.setProxyHost(proxy.getHostName());
                webRequest.setProxyPort(proxy.getPort());
            }
            //timeout
            this.webClient.getOptions().setTimeout(timeout);
            //request,response
            webClient.getPage(webRequest);
            HtmlPage page = webClient.getPage(request.getUrl());
            HttpResponse resp = new HttpResponse();
            WebResponse webResponse = page.getWebResponse();
            int status = webResponse.getStatusCode();
            resp.setStatus(status);
            if (status == 302 || status == 301) {
                String redirectUrl = webResponse.getResponseHeaderValue("Location");
                resp.setContent(UrlUtils.relative2Absolute(request.getUrl(), redirectUrl));
            } else if (status == 200) {
                String content = page.asXml();
                resp.setContent(content);
                InputStream inputStream = webResponse.getContentAsStream();
                ByteArrayOutputStream baos = new ByteArrayOutputStream();
                byte[] byteArr = new byte[1024];
                for (int index; (index = inputStream.read(byteArr)) != -1; ) {
                    baos.write(byteArr, 0, index);
                }
                byte[] buffer = baos.toByteArray();
                resp.setRaw(new ByteArrayInputStream(buffer));
                String contentType = webResponse.getContentType();
                resp.setContentType(contentType);
                String charset = getCharset(request.getCharset(), contentType);
                resp.setCharset(charset);
            } else {
                throw new DownloadException("ERROR : " + status);
            }
            System.out.println("Download page successfully: " + request.getUrl());
            return resp;
        } catch (Exception ex) {
            throw new DownloadException(ex);
        } finally {
            shutdown();
        }
    }

    public void shutdown() {
        webClient.close();
    }
}

